/*
 * Copyright (c) 2018-2022, NXOS Development Team
 * SPDX-License-Identifier: Apache-2.0
 *
 * Contains: Advanced Host Controller Interface for SATA
 * 
 * Change Logs:
 * Date           Author        Notes
 * 2022-08-13     JasonHu       first version
 */

#include <nxos_mini.h>

#ifdef CONFIG_NX_DRIVER_AHCI

#include "ahci.h"
#include <pci.h>
#include <dma_region.h>
#include <base/mutex.h>
#include <base/barrier.h>
#include <base/memory.h>
#include <base/thread.h>
#include <base/debug.h>
#include <base/driver.h>
#include <base/ioremap.h>
#include <base/irq.h>
#include <base/string.h>
#include <base/uaccess.h>
#include <base/block.h>
#include <base/page.h>
#define NX_LOG_NAME "AHCI"
#define NX_LOG_LEVEL NX_LOG_INFO
#include <base/log.h>
#include <base/malloc.h>

#define DRV_NAME "ahci sata driver"
#define DRV_VERSION "0.1"
#define DEV_NAME "sd"

/* memio info on the pci bar 5 */
#define PCI_AHCI_MEMIO_BAR   5

#define LOWER32(a) (NX_U32)((a) & 0xffffffff)
#define LOWER8(a) (NX_U8)((a) & 0xff)
#define HIGHER8(a) (NX_U8)(((a) >> 8)  & 0xff)

/* maxim ports we support */
#define DRV_AHCI_PORT_NR    32

struct HalDeviceExtension
{
    NX_U64 sector_count;   /* sectors in this disk. */
    NX_U8 type;    /* AHCI device type */
    NX_U8 port;    /* port for each device. */
    NX_U32 slots;  /* solts for device read/write transfer bits */
    NX_Mutex lock;   /* lock for disk read/write */
    void *fis_vaddr;
    void *clb_vaddr;
    NX_DmaRegion clb_dma;
    NX_DmaRegion fis_dma;
    void *cmd_hdrs[HBA_COMMAND_HEADER_NUM];             /* command header */
    NX_DmaRegion cmd_hdrs_dmas[HBA_COMMAND_HEADER_NUM];  /* command header dma */
};
typedef struct HalDeviceExtension HalDeviceExtension;

NX_PRIVATE struct hba_memory *g_hba_base; /* hba memory io base addr */

NX_PRIVATE NX_Error ahci_create_device(NX_Driver *driver, HalDeviceExtension *extension);

NX_PRIVATE NX_U32 ahci_flush_commands(struct hba_port *port)
{
    /* the commands may not take effect until the command
     * register is read again by software, because reasons.
     */
    NX_MemoryBarrierWrite();
    volatile NX_U32 c = port->command;
    NX_MemoryBarrier();
    return c;
}

NX_PRIVATE void ahci_stop_port_command_engine(struct hba_port *port)
{
    NX_MemoryBarrierWrite();
    port->command &= ~HBA_PxCMD_ST;
    NX_MemoryBarrierWrite();
    port->command &= ~HBA_PxCMD_FRE;
    NX_MemoryBarrier();
    while((port->command & HBA_PxCMD_CR) || (port->command & HBA_PxCMD_FR));
}

NX_PRIVATE void ahci_start_port_command_engine(struct hba_port *port)
{
    NX_MemoryBarrier();
    while(port->command & HBA_PxCMD_CR);
    NX_MemoryBarrierWrite();
    port->command |= HBA_PxCMD_FRE;
    NX_MemoryBarrierWrite();
    port->command |= HBA_PxCMD_ST;
    ahci_flush_commands((struct hba_port *)port);
}

NX_PRIVATE struct hba_command_header *ahci_initialize_command_header(HalDeviceExtension *dev, struct hba_memory *abar,
                                                                 struct hba_port *port, int slot, int write,
                                                                 int atapi, int prd_entries, int fis_len)
{
    struct hba_command_header *hdr = (struct hba_command_header *)dev->clb_vaddr;
    hdr += slot;
    hdr->write = write ? 1 : 0;
    hdr->prdb_count = 0;
    hdr->atapi=atapi ? 1 : 0;
    hdr->fis_length = fis_len;
    hdr->prdt_len = prd_entries;
    hdr->prefetchable = 0;
    hdr->bist = 0;
    hdr->pmport = 0;
    hdr->reset = 0;
    return hdr;
}

NX_PRIVATE struct fis_reg_host_to_device *ahci_initialize_fis_host_to_device(HalDeviceExtension *dev, struct hba_memory *abar,
                                                                         struct hba_port *port, int slot, int cmdctl, int ata_command)
{
    struct hba_command_table *tbl = (struct hba_command_table *)(dev->cmd_hdrs[slot]);
    struct fis_reg_host_to_device *fis = (struct fis_reg_host_to_device *)(tbl->command_fis);

    NX_MemZero(fis, sizeof(*fis));
    fis->fis_type = FIS_TYPE_REG_H2D;
    fis->command = ata_command;
    fis->c = cmdctl ? 1 : 0;
    return fis;
}

NX_PRIVATE void ahci_send_command(struct hba_port *port, int slot)
{
    port->interrupt_status = ~0;
    port->command_issue = (1 << slot);
    ahci_flush_commands(port);
}

NX_PRIVATE int ahci_write_prdt(HalDeviceExtension *dev, struct hba_memory *abar, struct hba_port *port,
                    int slot, int offset, int length, NX_Addr virt_buffer)
{
    int num_entries = ((length - 1) / PRDT_MAX_COUNT) + 1;
    struct hba_command_table *tbl = (struct hba_command_table *)(dev->cmd_hdrs[slot]);
    int i;
    struct hba_prdt_entry *prd;

    for(i = 0; i < num_entries - 1; i++)
    {
        NX_Addr phys_buffer;
        phys_buffer = NX_Virt2Phy(virt_buffer);
        prd = &tbl->prdt_entries[i + offset];
        prd->byte_count = PRDT_MAX_COUNT - 1;
        prd->data_base_l = LOWER32(phys_buffer);
        prd->data_base_h = 0;
        prd->interrupt_on_complete = 0;

        length -= PRDT_MAX_COUNT;
        virt_buffer += PRDT_MAX_COUNT;
    }

    NX_Addr phys_buffer;
    phys_buffer = NX_Virt2Phy(virt_buffer);
    prd = &tbl->prdt_entries[i + offset];
    prd->byte_count = length - 1;
    prd->data_base_l = LOWER32(phys_buffer);
    prd->data_base_h = 0;
    prd->interrupt_on_complete = 0;
    return num_entries;
}

NX_PRIVATE void ahci_reset_device(struct hba_memory *abar, struct hba_port *port, HalDeviceExtension *dev)
{
    NX_LOG_D("device port %d: sending COMRESET and reinitializing", dev->port);
    ahci_stop_port_command_engine(port);
    port->sata_error = ~0;
    /* power on, spin up */
    port->command |= 2;
    port->command |= 4;
    ahci_flush_commands(port);
    NX_MemoryBarrier();

    /* initialize state */
    port->interrupt_status = ~0; /* clear pending interrupts */
    port->interrupt_enable = AHCI_DEFAULT_INT; /* we want some interrupts */
    port->command &= ~((1 << 27) | (1 << 26)); /* clear some bits */
    port->sata_control |= 1;
    NX_MemoryBarrier();
    port->sata_control |= (~1);
    NX_MemoryBarrier();
    port->interrupt_status = ~0; /* clear pending interrupts */
    port->interrupt_enable = AHCI_DEFAULT_INT; /* we want some interrupts */
    NX_MemoryBarrier();
    ahci_start_port_command_engine(port);
    dev->slots = 0;
    port->sata_error = ~0;
}

NX_PRIVATE NX_Error ahci_port_dma_data_transfer(HalDeviceExtension *dev, struct hba_memory *abar, struct hba_port *port,
                                            int slot, int write, NX_Addr virt_buffer, int sectors, NX_U64 lba)
{
    struct fis_reg_host_to_device *fis;
    int timeout;
    int fis_len = sizeof(struct fis_reg_host_to_device) / 4;

    int ne = ahci_write_prdt(dev, abar, port, slot, 0, ATA_SECTOR_SIZE * sectors, virt_buffer);
    ahci_initialize_command_header(dev, abar, port, slot, write, 0, ne, fis_len);
    fis = ahci_initialize_fis_host_to_device(dev, abar, port, slot, 1, write ? ATA_CMD_WRITE_DMA_EX : ATA_CMD_READ_DMA_EX);

    fis->device = 1 << 6;
    fis->count_l = LOWER8(sectors);
    fis->count_h = HIGHER8(sectors);

    fis->lba0 = (unsigned char)( lba        & 0xFF);
    fis->lba1 = (unsigned char)((lba >> 8)  & 0xFF);
    fis->lba2 = (unsigned char)((lba >> 16) & 0xFF);
    fis->lba3 = (unsigned char)((lba >> 24) & 0xFF);
    fis->lba4 = (unsigned char)((lba >> 32) & 0xFF);
    fis->lba5 = (unsigned char)((lba >> 40) & 0xFF);
    port->sata_error = ~0;

    timeout = ATA_TFD_TIMEOUT;
    while ((port->task_file_data & (ATA_DEV_BUSY | ATA_DEV_DRQ)) && --timeout)
    {
        NX_ThreadYield();
    }
    if(!timeout)
    {
        goto port_hung;
    }

    port->sata_error = ~0;
    ahci_send_command(port, slot);
    timeout = ATA_TFD_TIMEOUT;
    while ((port->task_file_data & (ATA_DEV_BUSY | ATA_DEV_DRQ)) && --timeout)
    {
        NX_ThreadYield();
    }
    if(!timeout)
    {
        goto port_hung;
    }

    timeout = AHCI_CMD_TIMEOUT;
    while(--timeout)
    {
        if(!((port->sata_active | port->command_issue) & (1 << slot)))
            break;
        NX_ThreadYield();
    }
    if(!timeout)
    {
        goto port_hung;
    }
    if(port->sata_error)
    {
        NX_LOG_E("device %d: ahci error", dev->port);
        goto error;
    }
    if(port->task_file_data & ATA_DEV_ERR)
    {
        NX_LOG_E("device %d: task file data error", dev->port);
        goto error;
    }
    return NX_EOK;
port_hung:
    NX_LOG_E("device %d: port hung", dev->port);
error:
    NX_LOG_E("device %d: tfd=%x, serr=%x",
            dev->port, port->task_file_data, port->sata_error);
    ahci_reset_device(abar, port, dev);
    return NX_ERROR;
}

NX_PRIVATE NX_Error ahci_device_identify(HalDeviceExtension *dev, struct hba_memory *abar, struct hba_port *port)
{
    int fis_len = sizeof(struct fis_reg_host_to_device) / 4;
    NX_DmaRegion dma;
    dma.size = 0x1000;
    dma.alignment = 0x1000;
    NX_ASSERT(NX_DmaAllocBuffer(&dma) == NX_EOK);

    ahci_write_prdt(dev, abar, port, 0, 0, 512, (NX_Addr)dma.vaddr);
    ahci_initialize_command_header(dev, abar, port, 0, 0, 0, 1, fis_len);
    ahci_initialize_fis_host_to_device(dev, abar, port, 0, 1, ATA_CMD_IDENTIFY);

    int timeout = ATA_TFD_TIMEOUT;
    port->sata_error = ~0;
    while ((port->task_file_data & (ATA_DEV_BUSY | ATA_DEV_DRQ)) && --timeout);
    if(!timeout )
    {
        NX_LOG_E("device %d: identify 1: port hung", dev->port);
        NX_LOG_E("device %d: identify 1: tfd=%x, serr=%x",
                dev->port, port->task_file_data, port->sata_error);
        NX_DmaFreeBuffer(&dma);
        return NX_ETIMEOUT;
    }

    ahci_send_command(port, 0);

    timeout = AHCI_CMD_TIMEOUT;
    while(--timeout)
    {
        if(!((port->sata_active | port->command_issue) & 1))
            break;
    }
    if(!timeout)
    {
        NX_LOG_E("device %d: identify 2: port hung", dev->port);
        NX_LOG_E("device %d: identify 2: tfd=%x, serr=%x",
                dev->port, port->task_file_data, port->sata_error);
        NX_DmaFreeBuffer(&dma);
        return NX_ETIMEOUT;
    }

    struct ata_identify *identify = (struct ata_identify *) dma.vaddr;
    if (identify->lba48_addressable_sectors)
    {
        dev->sector_count = identify->lba48_addressable_sectors;
    }
    else
    {
        dev->sector_count = 0;
    }
    NX_LOG_I("port %d: num sectors=%d", dev->port, dev->sector_count);

    NX_DmaFreeBuffer(&dma);

    if (!dev->sector_count)
    {
        NX_LOG_E("device %d invalid sectors ZERO.", dev->port);
        return NX_EINVAL;
    }
    return NX_EOK;
}

NX_PRIVATE NX_U32 ahci_check_type(volatile struct hba_port *port)
{
    port->command &= ~1;
    while(port->command & (1 << 15));

    port->command &= ~(1 << 4);
    while(port->command & (1 << 14));

    NX_MemoryBarrierWrite();
    port->command |= 2;
    NX_MemoryBarrierWrite();

    NX_U32 s = port->sata_status;

    NX_U8 ipm, det;
    ipm = (s >> 8) & 0x0F;
    det = s & 0x0F;
    if(ipm != HBA_PORT_IPM_ACTIVE || det != HBA_PORT_DET_PRESENT)
    {
        return AHCI_DEV_NULL;
    }

    switch (port->signature)
    {
    case SATA_SIG_ATAPI:
        return AHCI_DEV_SATAPI;
    case SATA_SIG_SEMB:
        return AHCI_DEV_SEMB;
    case SATA_SIG_PM:
        return AHCI_DEV_PM;
    default:
        return AHCI_DEV_SATA;
    }
    return AHCI_DEV_SATA;
}

int ahci_initialize_device(HalDeviceExtension *dev, struct hba_memory *abar)
{
    struct hba_port *port = (struct hba_port *)&abar->ports[dev->port];
    ahci_stop_port_command_engine(port);
    port->sata_error = ~0;
    /* power on, spin up */
    port->command |= (2 | 4);
    ahci_flush_commands(port);
    NX_MemoryBarrier();

    /* initialize state */
    port->interrupt_status = ~0; /* clear pending interrupts */
    port->interrupt_enable = AHCI_DEFAULT_INT; /* we want some interrupts */

    port->command &= ~1;
    while(port->command & (1 << 15));

    port->command &= ~((1 << 27) | (1 << 26) | 1); /* clear some bits */
    ahci_flush_commands(port);

    /* start reset sata */
    port->sata_control |= 1;
    NX_MemoryBarrier();

    /* close DET, after init sata device done. */
    port->sata_control &= (~1);
    NX_MemoryBarrier();
    while(!(port->sata_status & 1));

    port->sata_error = ~0;
    port->command |= (1 << 28); /* set interface to active */
    while((port->sata_status >> 8) != 1);

    port->interrupt_status = ~0; /* clear pending interrupts */
    port->interrupt_enable = AHCI_DEFAULT_INT; /* we want some interrupts */

    NX_Addr clb_phys, fis_phys;
    dev->clb_dma.size = 0x2000;
    dev->clb_dma.alignment = 0x1000;
    dev->fis_dma.size = 0x1000;
    dev->fis_dma.alignment = 0x1000;

    NX_ASSERT(NX_DmaAllocBuffer(&dev->clb_dma) == NX_EOK);
    NX_ASSERT(NX_DmaAllocBuffer(&dev->fis_dma) == NX_EOK);
    dev->clb_vaddr = (void *)dev->clb_dma.vaddr;
    dev->fis_vaddr = (void *)dev->fis_dma.vaddr;
    clb_phys = dev->clb_dma.paddr;
    fis_phys = dev->fis_dma.paddr;

    dev->slots=0;
    struct hba_command_header *hdr = (struct hba_command_header *)dev->clb_vaddr;
    int i;
    for(i = 0; i < HBA_COMMAND_HEADER_NUM; i++)
    {
        dev->cmd_hdrs_dmas[i].size = 0x1000;
        dev->cmd_hdrs_dmas[i].alignment = 0x1000;

        NX_ASSERT(NX_DmaAllocBuffer(&dev->cmd_hdrs_dmas[i]) == NX_EOK);
        dev->cmd_hdrs[i] = (void *)dev->cmd_hdrs_dmas[i].vaddr;
        NX_MemZero(hdr, sizeof(*hdr));

        hdr->command_table_base_l = LOWER32(dev->cmd_hdrs_dmas[i].paddr);
        hdr->command_table_base_h = 0;

        hdr++;
    }

    port->command_list_base_l = LOWER32(clb_phys);
    port->command_list_base_h = 0;

    port->fis_base_l = LOWER32(fis_phys);
    port->fis_base_h = 0;
    ahci_start_port_command_engine(port);

    port->sata_error = ~0;
    return ahci_device_identify(dev, abar, port);
}

NX_PRIVATE NX_U32 ahci_probe_ports(NX_Driver *driver, struct hba_memory *abar)
{
    NX_U32 pi = abar->port_implemented;
    NX_LOG_D("ports implemented: %x", pi);
    int counts = 0; /* exist device count */
    int i = 0;
    HalDeviceExtension *extension;
    while (i < DRV_AHCI_PORT_NR)
    {
        if (pi & 1)
        {
            NX_U32 type = ahci_check_type(&abar->ports[i]);
            if (type == AHCI_DEV_SATA) { /* SATA device */
                NX_LOG_D("detected SATA device on port %d", i);

                extension = NX_MemAlloc(sizeof(HalDeviceExtension));
                if (extension == NX_NULL)
                {
                    NX_LOG_E("port %d alloc memory for extension failed!", i);
                    return counts;
                }
                extension->type = type;
                extension->port = i;
                NX_MutexInit(&extension->lock);
                if (ahci_initialize_device(extension, abar) == NX_EOK)\
                {
                    if (ahci_create_device(driver, extension) == NX_EOK) {
                        counts++;
                    }
                    else
                    {
                        NX_LOG_E("failed to create device %d, disabling port!", i);
                        NX_MemFree(extension);
                    }
                } else {
                    NX_LOG_E("failed to initialize device %d, disabling port.", i);
                }
            } else if(type == AHCI_DEV_SATAPI) { /* SATAPI device */
                NX_LOG_W("not support SATAPI device on port %d now!", i);
            } else if(type == AHCI_DEV_PM) { /* PM device */
                NX_LOG_W("not support Port multiplier on port %d now!", i);
            } else if(type == AHCI_DEV_SEMB) { /* SEMB device */
                NX_LOG_W("not support Enclosure management bridge on port %d now!", i);
            }
            /* do not deal other type now. */
        }
        i++;
        pi >>= 1;
    }
    return counts;
}

NX_PRIVATE int ahci_port_get_slot(HalDeviceExtension *dev)
{
    for(;;)
    {
        int i;
        NX_MutexLock(&dev->lock);
        for(i = 0; i < DRV_AHCI_PORT_NR; i++)
        {
            if(!(dev->slots & (1 << i)))
            {
                dev->slots |= (1 << i);
                NX_MutexUnlock(&dev->lock);
                return i;
            }
        }

        NX_MutexUnlock(&dev->lock);
        NX_ThreadYield();
    }
}

void ahci_port_put_slot(HalDeviceExtension *dev, int slot)
{
    NX_MutexLock(&dev->lock);
    dev->slots &= ~(1 << slot);
    NX_MutexUnlock(&dev->lock);
}

/* since a DMA transfer must write to contiguous physical RAM, we need to allocate
 * buffers that allow us to create PRDT entries that do not cross a page boundary.
 * That means that each PRDT entry can transfer a maximum of NX_PAGE_SIZE bytes (for
 * 0x1000 page size, that's 8 sectors). Thus, we allocate a buffer that is page aligned,
 * in a multiple of NX_PAGE_SIZE, so that the PRDT will write to contiguous physical ram
 * (the key here is that the buffer need not be contiguous across multiple PRDT entries).
 */
NX_PRIVATE NX_Size ahci_rw_multiple_do(HalDeviceExtension *dev, int rw, NX_U32 blk, unsigned char *out_buffer, int count)
{
    NX_U32 length = count * ATA_SECTOR_SIZE;
    NX_U32 end_blk = dev->sector_count;
    if (blk >= end_blk)
    {
        NX_LOG_E("ahci: lba %d out of range %d", blk, end_blk);
        return 0;
    }

    if((blk + count) > end_blk)
    {
        count = end_blk - blk;
    }
    if(!count)
    {
        return 0;
    }

    int num_pages = ((ATA_SECTOR_SIZE * (count - 1)) / NX_PAGE_SIZE) + 1;
    NX_ASSERT((length <= (unsigned)num_pages * 0x1000));
    NX_DmaRegion dma;
    dma.size = 0x1000 * num_pages;
    dma.alignment = 0x1000;
    NX_ASSERT(NX_DmaAllocBuffer(&dma) == NX_EOK);

    NX_Size num_read_blocks = count;
    struct hba_port *port = (struct hba_port *)&g_hba_base->ports[dev->port];
    if(rw == 1)
    {
        NX_MemCopy((void *)dma.vaddr, out_buffer, length);
    }

    int slot = ahci_port_get_slot(dev);
    if(ahci_port_dma_data_transfer(dev, g_hba_base, port, slot, rw == 1 ? 1 : 0, (NX_Addr)dma.vaddr, count, blk) != NX_EOK)
    {
        num_read_blocks = 0;
    }

    ahci_port_put_slot(dev, slot);

    if(rw == 0 && num_read_blocks)
    {
        NX_MemCopy(out_buffer, (void *)dma.vaddr, length);
    }

    NX_DmaFreeBuffer(&dma);
    return num_read_blocks;
}

/* and then since there is a maximum transfer amount because of the page size
 * limit, wrap the transfer function to allow for bigger transfers than that even.
 */
NX_PRIVATE NX_Size ahci_rw_multiple(HalDeviceExtension *dev, int rw, NX_U32 blk, unsigned char *out_buffer, int count)
{
    NX_LOG_D("trace ahci rw: %s, blk:%d, count:%d", rw == 0 ? "read" : "write", blk, count);
    int i = 0;
    NX_Size ret = 0;
    int c = count;
    for(i = 0; i < count; i += (PRDT_MAX_ENTRIES * PRDT_MAX_COUNT) / ATA_SECTOR_SIZE)
    {
        int n = (PRDT_MAX_ENTRIES * PRDT_MAX_COUNT) / ATA_SECTOR_SIZE;
        if(n > c)
        {
            n = c;
        }
        ret += ahci_rw_multiple_do(dev, rw, blk+i, out_buffer + ret, n);
        c -= n;
    }
    return ret;
}

NX_PRIVATE NX_PciDevice *ahci_get_pci_info(void)
{
    NX_PciDevice *ahci = NX_PciDeviceGet(0x1, 0x6);
    if(ahci == NX_NULL)
    {
        ahci = NX_PciDeviceGet(0x8086, 0x8c03);
    }
    if(ahci == NX_NULL)
    {
        ahci = NX_PciDeviceGet(0x8086, 0x2922);
    }
    if(ahci == NX_NULL)
    {
        return NX_NULL;
    }

    NX_LOG_D("device vendorID %x deviceID %x class code %x", ahci->vendor_id, ahci->device_id, ahci->class_code);

    NX_PciEnableBusMastering(ahci);

    g_hba_base = NX_IoRemap((void *) ahci->bars[PCI_AHCI_MEMIO_BAR].base_addr, ahci->bars[PCI_AHCI_MEMIO_BAR].length);
    if (g_hba_base == NX_NULL) {
        NX_LOG_D("device memio_remap on %x length %x failed!", ahci->bars[PCI_AHCI_MEMIO_BAR].base_addr, ahci->bars[PCI_AHCI_MEMIO_BAR].length);
        return NX_NULL;
    }
    NX_MmuFlush();

    NX_LOG_D("mapping memory iobase from paddr %x to vaddr %x", ahci->bars[PCI_AHCI_MEMIO_BAR].base_addr, g_hba_base);
    NX_LOG_D("using interrupt %d", ahci->irq_line);
    return ahci;
}

/*
 * position: block page address, not bytes address
 * buffer: read buffer addr
 * size  : how many blocks
 */

NX_PRIVATE NX_Error AHCI_Read(NX_Device *device, void *buf, NX_Offset off, NX_Size len, NX_Size *outLen)
{
    NX_Size readLen;
    NX_Size sectors = NX_DIV_ROUND_UP(len, ATA_SECTOR_SIZE);
    NX_LOG_D("trace ahci read: byte off: %p, byte len:%p sectors:%d", off, len, sectors);

    readLen = ahci_rw_multiple((HalDeviceExtension *)device->extension, 0, off / ATA_SECTOR_SIZE, (unsigned char *)buf, sectors);
    if (!readLen)
    {
        return NX_EIO;
    }
    readLen *= ATA_SECTOR_SIZE;
    NX_CopyToUser((char *)outLen, (char *)&readLen, sizeof(NX_Size));
    return NX_EOK;
}

/*
 * position: block page address, not bytes address
 * buffer: write buffer addr
 * size  : how many blocks
 */
NX_PRIVATE NX_Error AHCI_Write(NX_Device *device, void *buf, NX_Offset off, NX_Size len, NX_Size *outLen)
{
    NX_Size writeLen;
    NX_Size sectors = NX_DIV_ROUND_UP(len, ATA_SECTOR_SIZE);
    NX_LOG_D("trace ahci write: byte off: %p, byte len:%p sectors:%d", off, len, sectors);

    writeLen = ahci_rw_multiple((HalDeviceExtension *)device->extension, 1, off / ATA_SECTOR_SIZE, (unsigned char *)buf, sectors);
    if (!writeLen)
    {
        return NX_EIO;
    }
    writeLen *= ATA_SECTOR_SIZE;
    NX_CopyToUser((char *)outLen, (char *)&writeLen, sizeof(NX_Size));
    return NX_EOK;
}

NX_PRIVATE NX_Error AHCI_Control(NX_Device *device, NX_U32 cmd, void *arg)
{
    NX_ASSERT(device != NX_NULL);
    NX_ASSERT(device->extension != NX_NULL);
    NX_ASSERT(arg != NX_NULL);

    HalDeviceExtension *extension = (HalDeviceExtension *)device->extension;
    
    NX_Error err = NX_EOK;
    switch(cmd)
    {
    case NX_IO_BLOCK_INFO:
        {
            NX_IoBlockInfo *info = (NX_IoBlockInfo *)arg;
            info->capacity = extension->sector_count * ATA_SECTOR_SIZE;
            info->blockSize = ATA_SECTOR_SIZE;
            info->blockCount = extension->sector_count;
            
            NX_LOG_D("getgeome: capacity:%p, blockSize:%d, blockCount:%d",
                info->capacity, info->blockSize, info->blockCount);
            break;
        }
    default:
        err = NX_ERROR;
        break;
    }
    return err;
}

NX_PRIVATE NX_Error ahci_isr(NX_IRQ_Number vector, void *param)
{
    NX_Error err = NX_ENORES;
    int i;
    for (i = 0; i < 32; i++)
    {
        if (g_hba_base->interrupt_status & (1 << i))
        {
            NX_LOG_D("interrupt on port %d occured!", i);
            g_hba_base->ports[i].interrupt_status = ~0;
            g_hba_base->interrupt_status = (1 << i);
            ahci_flush_commands((struct hba_port *)&g_hba_base->ports[i]);
            err = NX_EOK;
            break;
        }
    }
    return err;
}

NX_PRIVATE void ahci_init_hba(struct hba_memory *abar)
{
    if(abar->ext_capabilities & 1)
    {
        /* request BIOS/OS ownership handoff */
        abar->bohc |= (1 << 1);
        while((abar->bohc & 1) || !(abar->bohc & (1<<1)));
    }
    /* enable the AHCI and reset it */
    abar->global_host_control |= HBA_GHC_AHCI_ENABLE;
    abar->global_host_control |= HBA_GHC_RESET;
    /* wait for reset to complete */
    while(abar->global_host_control & HBA_GHC_RESET);

    /* enable the AHCI and interrupts */
    abar->global_host_control |= HBA_GHC_AHCI_ENABLE;
    abar->global_host_control |= HBA_GHC_INTERRUPT_ENABLE;
    NX_MemoryBarrier();
    NX_LOG_D("caps: 0x%p ext: 0x%p ver:%x ctl: %x", abar->capability, abar->ext_capabilities, abar->version, abar->global_host_control);
}

NX_PRIVATE NX_Error ahci_create_device(NX_Driver *driver, HalDeviceExtension *extension)
{
    NX_PRIVATE int ahci_next_device = 0;   /* first is sd0 */
    NX_Device *device;

    char devname[8] = {0};
    NX_SNPrintf(devname, 8, "%s%c", DEV_NAME, '0' + ahci_next_device);
    ahci_next_device++;

    if (NX_DriverAttachDevice(driver, devname, &device) != NX_EOK)
    {
        NX_LOG_E("attach device %s failed!", DEV_NAME);
        return NX_ENOMEM;
    }

    device->extension = extension;
    
    return NX_EOK;
}

NX_PRIVATE int ahci_init(NX_Driver *driver)
{
    /* 1. get pci info */
    NX_PciDevice *ahci_pci = ahci_get_pci_info();
    if(ahci_pci == NX_NULL)
    {
        NX_LOG_E("no AHCI controllers present!");
        return NX_ERROR;
    }

    /* 2. install intr */
    if (NX_IRQ_Bind(ahci_pci->irq_line, ahci_isr, ahci_pci, "ahci", NX_IRQ_FLAG_SHARED) != NX_EOK)
    {
        NX_IoUnmap(g_hba_base);
        return NX_ERROR;
    }
    NX_IRQ_Unmask(ahci_pci->irq_line);

    /* 3. init ahci device */
    ahci_init_hba(g_hba_base);
    if (!ahci_probe_ports(driver, g_hba_base))
    {
        NX_LOG_E("initializing ahci driver failed!.");
        NX_IRQ_Mask(ahci_pci->irq_line);
        NX_IoUnmap(g_hba_base);
        return NX_ERROR;
    }
    NX_LOG_I("disk driver init done!");
    return NX_EOK;
}

NX_PRIVATE NX_DriverOps AHCI_DriverOps = {
    .read       = AHCI_Read,
    .write      = AHCI_Write,
    .control    = AHCI_Control,
};

NX_PRIVATE void AHCI_DriverInit(void)
{
    NX_Driver *driver = NX_DriverCreate(DRV_NAME, NX_DEVICE_TYPE_BLOCK, 0, &AHCI_DriverOps);
    if (driver == NX_NULL)
    {
        NX_LOG_E("create driver failed!");
        return;
    }

    /* call init */
    if (ahci_init(driver) != NX_EOK)
    {
        return;
    }

    if (NX_DriverRegister(driver) != NX_EOK)
    {
        NX_LOG_E("register driver %s failed!", DRV_NAME);
        NX_DriverCleanup(DRV_NAME);
        return;
    }
}

NX_PRIVATE void AHCI_DriverExit(void)
{
    NX_IoUnmap(g_hba_base);
    NX_DriverCleanup(DRV_NAME);
}

NX_DRV_INIT(AHCI_DriverInit);
NX_DRV_EXIT(AHCI_DriverExit);

#endif /* CONFIG_NX_DRIVER_AHCI */
